#include "../DWMRICore/DTI.h"
#include "../DWMRICore/HOT.h"
#include "../DWMRICore/Scalar.h"

char *dti_pathname = NULL;
char *hot_pathname = NULL;
char *mask_pathname = NULL;
char *dst_pathname = NULL;

CDTI *pDTI = NULL;
CHOT *pHOT = NULL;
CScalar *pMASK = NULL;
uint w = 0;
uint h = 0;
uint d = 0;
uint n = 0;

void PrintUsage()
{
	printf("Usage: MappingFibers -dti dti_file -hot hot_file -mask mask_file -dst dst_file\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			dti_pathname = argv[++i];
		} else if (strcmp(argv[i], "-hot") == 0) {
			hot_pathname = argv[++i];
		} else if (strcmp(argv[i], "-mask") == 0) {
            mask_pathname = argv[++i];
        } else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else {
			return -1;
		}
	}

	if (dti_pathname == NULL || hot_pathname == NULL || mask_pathname == NULL || dst_pathname == NULL)
		return -2;

	return 0;
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(dti_pathname);
	pHOT = new CHOT();
	pHOT->OpenHOTFile(hot_pathname);
	pMASK = new CScalar();
	pMASK->OpenScalarFile(mask_pathname);

	w = pDTI->GetWidth();
	h = pDTI->GetHeight();
	d = pDTI->GetDepth();
	n = pHOT->GetN();

	float *vectorField = new float[6*w*h*d];
	memset(vectorField, 0, sizeof(float)*6*w*h*d);
	for (uint i = 0; i < w * h * d; ++i) {
		float mask = pMASK->GetScalarFast(i);
		if (mask < 0.5f) {
			;
		} else if (mask < 1.5f) {
			float *dti = new float[7];
			pDTI->GetDTIFast(i, dti);
			CVector3F dir;
			pDTI->DTI2Direction(dti, &dir);
			vectorField[i*6+0] = dir.m_x;
			vectorField[i*6+1] = dir.m_y;
			vectorField[i*6+2] = dir.m_z;
			delete[] dti;
		} else {
			float *hot = new float[n];
			pHOT->GetHOTFast(i, hot);
			CVector3F dir[3];
			pHOT->HOT2Directions_Two(hot, dir);
			vectorField[i*6+0] = dir[0].m_x;
			vectorField[i*6+1] = dir[0].m_y;
			vectorField[i*6+2] = dir[0].m_z;
			vectorField[i*6+3] = dir[1].m_x;
			vectorField[i*6+4] = dir[1].m_y;
			vectorField[i*6+5] = dir[1].m_z;
			delete[] hot;
		}
	}

	/* save result */
	FILE *fp = fopen(dst_pathname, "wb");
	fwrite(vectorField, sizeof(float), 6*w*h*d, fp);
	fclose(fp);

	return 0;
}